#include <iostream>
#include <string>
#include <vector>
#include <filesystem>
#include <iomanip>
#include <cmath>
#include <thread>
#include <chrono>
#include <omp.h>
#include "DataTypes.h"
#include "FileIO.h"
#include "Optimizer.h"
#include "CoreLogic.h"
#include "VisualizationExporter.h"

int main() {
    unsigned int max_threads = std::thread::hardware_concurrency();
    unsigned int num_threads = (max_threads > 1) ? max_threads - 1 : 1;
    omp_set_num_threads(num_threads);
    std::cout << "--- Parallel computing enabled with " << num_threads << " threads (Max hardware threads: " << max_threads << ") ---" << std::endl;

    // Define all cases to run
    std::vector<std::string> case_names = {
        "Conv_Case0",
        "Conv_Case1",
        "Matmul_Case0",
        "Matmul_Case1",
        "FlashAttention_Case0",
        "FlashAttention_Case1",
    };

    std::cout << "=== Running Multi-Case Optimization Suite ===" << std::endl;
    std::cout << "Total cases to process: " << case_names.size() << std::endl;

    for (size_t case_idx = 0; case_idx < case_names.size(); ++case_idx) {
        const std::string& case_name = case_names[case_idx];
        
        std::cout << "\n" << std::string(80, '=') << std::endl;
        std::cout << "Processing Case " << (case_idx + 1) << "/" << case_names.size() 
                  << ": " << case_name << std::endl;
        std::cout << std::string(80, '=') << std::endl;

        try {
            // Hardware Capabilities (Caps) - same for all cases
            Caps caps;
            caps.l1 = 4096;   // L1 cache capacity
            caps.ub = 1024;   // UB cache capacity
            caps.l0a = 256;   // L0A cache capacity
            caps.l0b = 256;   // L0B cache capacity
            caps.l0c = 512;   // L0C cache capacity

            // Initial Control Parameters (Params)
            Params params;
            params.min_farthest_use_threshold = 500;  // Restore default used in full run
            params.w_use = 0.5;
            params.w_vol = 0.3;
            params.w_vic = 0.2;
            params.alloc_strategy = "BestFit";

            // Options for NSGA-II genetic algorithm - Adjusted for debugging
            Options options;
            options.generations = 500;  // Restore full run generations
            options.population_size = 10;  // Restore typical population size
            options.batch_size = 10;
            options.schedule_move_prob = 0.8;
            options.threshold_range = { 100.0, 2000.0 };  // Restore original range
            
            std::cout << "Configuration: " << options.generations << " generations, " 
                      << options.population_size << " population size" << std::endl;

            // File Loading with size checks
            const std::string nodes_file = std::string(DATA_DIR) + "/" + case_name + "_Nodes.csv";
            const std::string edges_file = std::string(DATA_DIR) + "/" + case_name + "_Edges.csv";
            const std::string schedule_file = std::string(DATA_DIR) + "/" + case_name + "_final_schedule_optimized.txt";
            
            std::cout << "File sizes:" << std::endl;
            std::cout << "  Nodes: " << (std::filesystem::file_size(nodes_file) / 1024.0) << " KB" << std::endl;
            std::cout << "  Edges: " << (std::filesystem::file_size(edges_file) / 1024.0) << " KB" << std::endl;
            std::cout << "  Schedule: " << (std::filesystem::file_size(schedule_file) / 1024.0) << " KB" << std::endl;
            
            // Debug info - estimated complexity
            std::cout << "Estimated nodes from file size: " << (std::filesystem::file_size(nodes_file) / 50) << std::endl;

            // Check if all required files exist
            if (!std::filesystem::exists(nodes_file)) {
                std::cerr << "ERROR: Missing nodes file: " << nodes_file << std::endl;
                continue;
            }
            if (!std::filesystem::exists(edges_file)) {
                std::cerr << "ERROR: Missing edges file: " << edges_file << std::endl;
                continue;
            }
            if (!std::filesystem::exists(schedule_file)) {
                std::cerr << "ERROR: Missing schedule file: " << schedule_file << std::endl;
                continue;
            }

            Graph graph;
            CodeMap code_map;
            if (!FileIO::loadGraph(nodes_file, edges_file, graph, code_map)) {
                std::cerr << "ERROR: Failed to load graph for " << case_name << std::endl;
                continue;
            }
            std::cout << "Graph loaded: " << graph.num_nodes << " nodes." << std::endl;

            std::vector<int> p1_schedule;
            try {
                p1_schedule = FileIO::loadScheduleFromFile(schedule_file);
            }
            catch (const std::runtime_error& e) {
                std::cerr << "ERROR: Failed to load schedule for " << case_name << ": " << e.what() << std::endl;
                continue;
            }
            std::cout << "Schedule loaded, optimizing " << p1_schedule.size() << " nodes." << std::endl;

            // Initialize Visualization Exporter for this case
            std::string viz_output_dir = "output/visualizations/" + case_name;
            VisualizationExporter viz_exporter(viz_output_dir, case_name);
            std::cout << "Visualization exporter initialized for " << case_name << std::endl;

            // Execute Optimization with error handling
            std::vector<Individual> pareto_front;
            auto start_time = std::chrono::high_resolution_clock::now();
            
            try {
                std::cout << "Graph complexity: " << graph.nodes.size() << " nodes" << std::endl;
                std::cout << "Starting optimization..." << std::endl;
                
                OptimizerNSGA optimizer(p1_schedule, graph, code_map, caps, params, options, &viz_exporter);
                pareto_front = optimizer.run();
                
                auto end_time = std::chrono::high_resolution_clock::now();
                auto duration = std::chrono::duration_cast<std::chrono::seconds>(end_time - start_time);
                std::cout << "Optimization completed in " << duration.count() << " seconds" << std::endl;
                
            } catch (const std::exception& e) {
                std::cerr << "ERROR processing " << case_name << ": Optimization failed: " << e.what() << std::endl;
                continue;
            } catch (...) {
                std::cerr << "ERROR processing " << case_name << ": Unknown optimization error" << std::endl;
                continue;
            }

            // Results Analysis and Saving
            std::cout << "Found " << pareto_front.size() << " Pareto optimal solutions" << std::endl;
            if (!pareto_front.empty()) {
                std::filesystem::path case_output_dir = std::string("output/Problem3_CPP/") + case_name;
                std::filesystem::create_directories(case_output_dir);

                std::vector<double> all_times, all_spills;
                for (int i = 0; i < pareto_front.size(); ++i) {
                    const auto& sol = pareto_front[i];
                    all_times.push_back(sol.time);
                    all_spills.push_back(sol.spill);
                    
                    // Collect pareto solution data
                    ParetoSolution viz_sol;
                    viz_sol.solution_id = i + 1;
                    viz_sol.time = sol.time;
                    viz_sol.spill = sol.spill;
                    viz_sol.strategy_type = "unknown"; // Can be enhanced to add strategy info
                    viz_sol.params = sol.params;
                    viz_exporter.addParetoSolution(viz_sol);
                }

                auto min_time_it = std::min_element(all_times.begin(), all_times.end());
                auto min_spill_it = std::min_element(all_spills.begin(), all_spills.end());
                int min_time_idx = std::distance(all_times.begin(), min_time_it);
                int min_spill_idx = std::distance(all_spills.begin(), min_spill_it);

                double min_t = *min_time_it;
                double range_t = *std::max_element(all_times.begin(), all_times.end()) - min_t;
                double min_s = *min_spill_it;
                double range_s = *std::max_element(all_spills.begin(), all_spills.end()) - min_s;

                int min_dist_idx = -1;
                double min_dist = std::numeric_limits<double>::infinity();
                for (int i = 0; i < pareto_front.size(); ++i) {
                    double norm_t = (range_t > 0) ? (all_times[i] - min_t) / range_t : 0;
                    double norm_s = (range_s > 0) ? (all_spills[i] - min_s) / range_s : 0;
                    double dist = std::sqrt(norm_t * norm_t + norm_s * norm_s);
                    if (dist < min_dist) {
                        min_dist = dist;
                        min_dist_idx = i;
                    }
                }

                std::cout << "Key solutions: Best Time (ID " << (min_time_idx+1) << "), "
                          << "Best Spill (ID " << (min_spill_idx+1) << "), "
                          << "Best Tradeoff (ID " << (min_dist_idx+1) << ")" << std::endl;

                // Save ALL solutions (same as original main.cpp)
                std::cout << "Validating and saving all optimal solutions" << std::endl;
                for (int i = 0; i < pareto_front.size(); ++i) {
                    const auto& current_solution_from_front = pareto_front[i];

                    std::string folder_name = std::string("Solution_") + (i < 9 ? "0" : "") + std::to_string(i + 1);
                    if (i == min_time_idx) folder_name += "_Best_Time";
                    if (i == min_spill_idx) folder_name += "_Best_Spill";
                    if (i == min_dist_idx) folder_name += "_Best_Tradeoff";

                    std::cout << "Processing solution #" << (i + 1) << " (" << folder_name << ")" << std::endl;

                    std::filesystem::path solution_dir = case_output_dir / folder_name;
                    std::filesystem::create_directories(solution_dir);

                    auto [solution_obj, final_graph] = CoreLogic::performAllocation(
                        current_solution_from_front.schedule, graph, code_map, caps, current_solution_from_front.params, &viz_exporter);

                    std::cout << "Validating solution for saving..." << std::endl;
                    bool is_valid = CoreLogic::validateSolution(solution_obj, final_graph, code_map, caps);

                    if (is_valid) {
                        std::cout << "  Validation PASSED. Saving TXT files..." << std::endl;
                        FileIO::writeScheduleToFile(solution_dir / (case_name + "_schedule.txt"), solution_obj.final_schedule);
                        FileIO::writeMemoryToFile(solution_dir / (case_name + "_memory.txt"), solution_obj.memory_offsets);
                        auto spill_log = FileIO::formatSpillLogForOutput(solution_obj.spill_map, final_graph);
                        FileIO::writeSpillToFile(solution_dir / (case_name + "_spill.txt"), spill_log);
                        std::cout << "  TXT results saved to: " << solution_dir << std::endl;
                    }
                    else {
                        std::cerr << "!!! Validation FAILED: Solution #" << (i + 1) << " has issues. Files will not be saved. !!!" << std::endl;
                    }
                }
            }
            
            // Export Visualization Data
            viz_exporter.exportAllData();
            std::cout << "Visualization data exported for " << case_name << std::endl;

        }
        catch (const std::exception& e) {
            std::cerr << "ERROR processing " << case_name << ": " << e.what() << std::endl;
            continue;
        }

        std::cout << "Case " << case_name << " completed successfully!" << std::endl;
    }

    std::cout << "\n" << std::string(80, '=') << std::endl;
    std::cout << "=== Multi-Case Optimization Suite Completed ===" << std::endl;
    std::cout << "Check output/Problem3_CPP/ for solution files" << std::endl;
    std::cout << "Check output/visualizations/ for visualization data" << std::endl;
    std::cout << std::string(80, '=') << std::endl;

    return 0;
}
